import os
import shutil

# 源文件夹路径
source_dir = 'deepglobe'

# 目标文件夹路径
target_dir = 'deepglobe_reorganized'

# 创建目标目录结构
os.makedirs(os.path.join(target_dir, 'gt'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'images'), exist_ok=True)

# 复制train中的mask和sat图像
train_dir = os.path.join(source_dir, 'train')
for filename in os.listdir(train_dir):
    if filename.endswith('_mask.png'):
        # 复制mask图像
        shutil.copy(os.path.join(train_dir, filename), os.path.join(target_dir, 'gt', filename))
    # elif filename.endswith('_sat.jpg'):
    #     # 复制sat图像
    #     shutil.copy(os.path.join(train_dir, filename), os.path.join(target_dir, 'images', filename))

# 注意：根据您的描述，valid和road_test_sat文件夹只包含sat图像，因此我们不需要从这些文件夹复制mask图像。
# 如果您决定将valid和road_test_sat文件夹中的sat图像也包含在训练/验证过程中，请取消注释并适当修改下面的代码。

# 复制valid中的sat图像
# valid_dir = os.path.join(source_dir, 'valid')
# for filename in os.listdir(valid_dir):
#     if filename.endswith('_sat.jpg'):
#         shutil.copy(os.path.join(valid_dir, filename), os.path.join(target_dir, 'images', filename))

# 复制road_test_sat中的sat图像
# road_test_sat_dir = os.path.join(source_dir, 'road_test_sat')
# for filename in os.listdir(road_test_sat_dir):
#     if filename.endswith('_sat.jpg'):
#         shutil.copy(os.path.join(road_test_sat_dir, filename), os.path.join(target_dir, 'images', filename))

print("数据集重组完成！")
